import cv2
import os
import numpy as np

# 根据文件名，以灰度格式读入图片，返回一维向量（ndarray）
def img2vector(imgfilename):
    img = cv2.imread(imgfilename, cv2.IMREAD_GRAYSCALE )
    rows, columns = img.shape
    img = img.reshape(rows*columns)
    return img

# 获取训练目录下的所有子目录
training_dir = '../training'
sub_dir_and_files = os.listdir(training_dir)
sub_dirs = []
for x in sub_dir_and_files:
    if os.path.isdir(training_dir + '\\' + x):
        sub_dirs.append(x)

# 计算训练目录下所有装备的图像数量和
N = 0
for subdir in sub_dirs:
    N += len(os.listdir(training_dir + '\\' + subdir))

# 初始化训练图像数据矩阵（N行，128*128列）和（相对应的）装备向量（长度为N）
training_image_matrix = np.zeros((N, 128*128)) # 每个训练图像对应一行所有像素值
training_equipment_vector = ['']*N # 与同下标位置训练图像相对应的装备名称

# 将三个子目录下所有图像文件读入图像数据矩阵并标识装备类别名称
i = 0 # 记录当前下标位置
for subdir in sub_dirs:
    image_files = os.listdir(training_dir + '\\' + subdir)
    for image in image_files:
        # 图像不在当前目录，因此需要加上目录前缀
        v = img2vector(training_dir+'\\'+subdir+'\\'+image)
        training_image_matrix[i] = v
        training_equipment_vector[i] = subdir
        i += 1
        
#print(training_image_matrix.shape)
#print(training_equipment_vector)

# 计算两个长度相同的向量之间的欧几里得距离（euclidean distance）：
def distance(vector1, vector2):
    diff = vector1-vector2   # 计算差向量
    double_diff = diff**2     # 计算差平方向量
    sum_of_double_diff = double_diff.sum()  # 计算两向量差平方和
    dist = np.sqrt(sum_of_double_diff)
    return dist

# 识别目标的类别，返回类别名称
# training_image_matrix是训练集所有图像的数据矩阵，class_vector是各图像类别向量
def classify(img, training_image_matrix, class_vector, k):
    N = len(training_image_matrix)  # 训练集图像个数
    dist = np.zeros((N))  # 距离向量初始化为全0
    # 计算待识别图像与训练集各个图像之间的欧几里得距离
    for i in range(N):
        dist[i] = distance(img, training_image_matrix[i])
    # 对距离进行排序（间接排序，否则类别名称会对不上）
    sorted_indicies = dist.argsort()  # 间接排序，返回从小到大排序后的元素下标
    # 统计前k个距离最小的匹配对象及其匹配次数
    match_count = {} # 用字典方式，便于数据结对
    for i in range(k):  
        match_class = class_vector[sorted_indicies[i]]
        match_count[match_class] = match_count.get(match_class,0) + 1  # 找不到则缺省为0
    
    match_count_in_order = sorted(match_count.items(), key = lambda item: item[1], reverse = True)  # 按照匹配次数从高到底排序，转化为二维列表了
    
    recognized = match_count_in_order[0][0]  # 匹配次数最高的类别名称
    return recognized

# 通常采用测试用例的方式进行测试
testcases = ['bmp-2-1.jpg', 'btr-70-1.jpg', 't-72-1.jpg', 't-72-3.jpg', 't-72-25.jpg']
for img in testcases:
    imgvector = img2vector(img)
    recognized = classify(imgvector, training_image_matrix, training_equipment_vector, 3)
    print('图像%12s目标识别结果为：'%img, recognized)

# 获取测试目录下的所有子目录
testing_dir = '../testing'
sub_dir_and_files = os.listdir(training_dir)
sub_dirs = []
for x in sub_dir_and_files:
    if os.path.isdir(testing_dir + '\\' + x):
        sub_dirs.append(x)


# 计算测试目录下所有装备图像的数量和
N = 0
for subdir in sub_dirs:
    N += len(os.listdir(testing_dir + '\\' + subdir))

# 初始化测试图像数据矩阵（N行，128*128列）和（相对应的）装备向量（长度为N）
testing_image_matrix = np.zeros((N, 128*128)) # 每个测试图像对应一行所有像素值
testing_equipment_vector = [""]*N # 与同下标位置测试图像相对应的装备名称

# 将三个子目录下所有图像文件读入图像数据矩阵并标识装备类别名称
i = 0 # 记录当前下标位置
for subdir in sub_dirs:
    image_files = os.listdir(testing_dir + "\\" + subdir)
    for image in image_files:
        # 图像不在当前目录，因此需要加上目录前缀
        v = img2vector(testing_dir+"\\"+subdir+"\\"+image)
        testing_image_matrix[i] = v
        testing_equipment_vector[i] = subdir
        i += 1

correct_count = 0
for i in range(N):
    k = 3
    recognized = classify(testing_image_matrix[i], training_image_matrix, training_equipment_vector, k)
    if recognized == testing_equipment_vector[i]:
        print('  正确：', recognized)
        correct_count +=1
    else:
        print('不正确：', testing_equipment_vector[i], '误识别为', recognized)
print('%d个测试对象判定正确%d个，正确率为：%.1f%%' %(N, correct_count, correct_count/N*100))

### 绘制正确率柱形图
##import matplotlib.pyplot as plt
##from pylab import mpl
##mpl.rcParams['font.sans-serif']= ['SimHei']
##x = [0.3, 1, 1.7]
##y = [93.7, 85.9, 86.2] # 识别正确率（BMP-2, BTR-70, T-72）
##plt.bar(x, y, width=0.4, color='b')
##xlabels = ['BMP-2', 'BTR-70', 'T-72']
##plt.xticks([0.5, 1.2, 1.9], xlabels)
##plt.title('各种装备的识别正确率（%）')
##plt.axis([0,2.3,0,100])
##plt.grid('on')
##plt.show()
